#include "core/main.h"

namespace GridWorldConstants
{
    static const int gridSize = 5;

    // Size of the player sprite (in pixels)
    static const int playerSize = 16;

    // The player starts on this square.
    static const int startX = 0;
    static const int startY = 0;

    // Stepping on this square gives a reward.
    static const int targetX = gridSize - 1;
    static const int targetY = gridSize - 1;
    static const int reward = 100;
    
    // Stepping on this square gives a penalty.
    static const int badX = 2;
    static const int badY = 2;
    static const int penalty = 25;
}

void GameGridWorld::Init()
{
    _spritePlayer = _collection.AddSprite(Bitmap::LoadFile("../assets/gridWorld/player.png"), "player");

    _collection.AddFontSprites();

    _observedIntNames.push_back("Reset");
    _observedIntNames.push_back("Score");
    _observedIntNames.push_back("Frame score");
    _observedIntNames.push_back("Player x");
    _observedIntNames.push_back("Player y");
    _observedIntNames.push_back("Frames since reset");
    _observedIntNames.push_back("dummy");

    _activeFrame.observedInts.resize(ObservedIntNames().size());
}

void GameGridWorld::NewGame()
{
    _state.score = 0;
    Reset();
}

void GameGridWorld::Reset()
{
    _state.x = GridWorldConstants::startX;
    _state.y = GridWorldConstants::startY;
    _state.framesSinceReset = 0;
}

void GameGridWorld::Step(ControllerState &controller)
{
    int frameScore = 0;

    if (controller.left()) {
        _state.x--;
    } else if (controller.right()) {
        _state.x++;
    } else if (controller.up()) {
        _state.y++;
    } else if (controller.down()) {
        _state.y--;
    }

    _state.x = Math::Bound(_state.x, 0, GridWorldConstants::gridSize - 1);
    _state.y = Math::Bound(_state.y, 0, GridWorldConstants::gridSize - 1);

    bool reset = false;
    if (_state.x == GridWorldConstants::targetX && _state.y == GridWorldConstants::targetY) {
        frameScore += GridWorldConstants::reward;
        reset = true;
    }

    if (!reset && _state.x == GridWorldConstants::badX && _state.y == GridWorldConstants::badY) {
        frameScore -= GridWorldConstants::penalty;
    }

    _state.score += frameScore;
    _state.framesSinceReset++;

    _activeFrame.observedInts[0] = reset ? 1 : 0;
    _activeFrame.observedInts[1] = _state.score;
    _activeFrame.observedInts[2] = frameScore;
    _activeFrame.observedInts[3] = _state.x;
    _activeFrame.observedInts[4] = _state.y;
    _activeFrame.observedInts[5] = _state.framesSinceReset;
    _activeFrame.observedInts[6] = 0; // dummy

    if (reset) Reset();

    MakeFrameSprites();
}

ControllerState GameGridWorld::SimpleAI()
{
    ControllerState controller;

    // Move randomly
    float num = rnd();
    if (num > 0.2f && num <= 0.4f) controller.KeyDown('a');
    else if (num > 0.4f && num <= 0.6f) controller.KeyDown('s');
    else if (num > 0.6f && num <= 0.8f) controller.KeyDown('d');
    else if (num > 0.8f) controller.KeyDown('w');

    return controller;
}

void GameGridWorld::MakeFrameSprites()
{
    _activeFrame.sprites.clear();

    // Draw the player
    int screenX = _state.x * GridWorldConstants::playerSize;
    int screenY = (GridWorldConstants::gridSize - _state.y - 1) * GridWorldConstants::playerSize;
    _activeFrame.sprites.push_back(SpriteLocation(Vec2i(screenX, screenY), _spritePlayer, 0));

    // Draw the score
    _collection.DrawText("SCORE", Vec2i(24, 1), _activeFrame.sprites, 0);
    _collection.DrawNumber(_state.score, 5, Vec2i(24, 2), _activeFrame.sprites, 0);
}